import cv2
import numpy as np
import matplotlib.pyplot as plt
import os


def detect_lanes(image_path):
    # 获取输入文件名（不含扩展名）
    filename = os.path.splitext(os.path.basename(image_path))[0]

    # 读取图像
    color_img = cv2.imread(image_path)
    gray_img = cv2.imread(image_path, 0)  # 直接读取灰度图

    # 形态学闭运算
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    gray_img = cv2.morphologyEx(gray_img, cv2.MORPH_CLOSE, kernel)

    # 二值化
    _, binary = cv2.threshold(gray_img, 125, 255, cv2.THRESH_BINARY)

    # Canny边缘检测
    canny = cv2.Canny(binary, 50, 125, 3)

    # 霍夫直线检测
    lines = cv2.HoughLinesP(
        canny,
        rho=1,
        theta=np.pi / 180,
        threshold=100,
        minLineLength=10,
        maxLineGap=50
    )

    # 绘制直线
    result = color_img.copy()
    if lines is not None:
        for line in lines:
            x1, y1, x2, y2 = line[0]
            # 计算角度
            dx = x2 - x1
            dy = y2 - y1
            angle = np.arctan2(float(dy), dx) * 180 / np.pi


            cv2.line(result, (x1, y1), (x2, y2), (255, 0, 0), 1)

    # 显示结果
    plt.figure(figsize=(15, 10))

    plt.subplot(221)
    plt.imshow(cv2.cvtColor(gray_img, cv2.COLOR_GRAY2RGB))
    plt.title('Gray Image')

    plt.subplot(222)
    plt.imshow(cv2.cvtColor(binary, cv2.COLOR_GRAY2RGB))
    plt.title('Binary Image')

    plt.subplot(223)
    plt.imshow(cv2.cvtColor(canny, cv2.COLOR_GRAY2RGB))
    plt.title('Canny Edge')

    plt.subplot(224)
    plt.imshow(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
    plt.title('Result')

    plt.tight_layout()
    plt.show()

    # 保存结果图像
    output_path = f"{filename}end.jpg"
    cv2.imwrite(output_path, result)
    print(f"结果已保存为: {output_path}")

    return result


if __name__ == "__main__":
    image_path = "road2.jpg"  # 输入图片路径
    result = detect_lanes(image_path)
